package com.example.aigc_education.ws;


import com.example.aigc_education.domain.KgChatConversation;
import com.example.aigc_education.domain.request.ChatAiContentRequest;
import com.example.aigc_education.integration.AgentServiceIntegration;
import com.example.aigc_education.integration.request.ChatMessageRequestV2;
import com.example.aigc_education.security.domain.LoginUser;
import com.example.aigc_education.service.KgChatConversationService;
import com.example.aigc_education.utils.*;
import lombok.extern.slf4j.Slf4j;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.security.Principal;
import java.util.Date;

@Slf4j
public class PlusWebSocketHandler extends AbstractWebSocketHandler {
    /**
     * 连接成功后
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        WebSocketSessionHolder.addSession(session.getId(), session);
    }

    /**
     * 处理发送来的文本消息
     *
     * @param session WebSocket会话
     * @param message 接收到的文本消息
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) {
        String sessionId = session.getId();
        String payload = message.getPayload();
        log.info("[receive message] sessionId: {}", sessionId);
        
        // 获取用户信息
        LoginUser user = extractLoginUser(session);
        if (user == null) {
            log.warn("WebSocket连接未获取到用户信息，sessionId: {}", sessionId);
            WebSocketUtils.sendMessage(session, "未授权的连接");
            return;
        }
        
        log.info("[receive message] WebSocket连接用户: {}", user.getUser().getUserId());

        try {
            ChatAiContentRequest userContentRequest = JsonUtils.fromJson(payload, ChatAiContentRequest.class);
            if (userContentRequest == null) {
                log.warn("无效的请求格式, sessionId: {}", sessionId);
                WebSocketUtils.sendMessage(session, "无效的请求格式");
                return;
            }
            if (StringUtils.isEmpty(userContentRequest.getConvId())) {
                // TODO: 2025/3/16 处理convId为null的情况，及创建新会话
                log.info("创建新会话, userId: {}", user.getUser().getUserId());

                // 创建新会话
                KgChatConversation conversation = new KgChatConversation();
                // 设置默认标题，可以使用用户消息的前N个字符或者固定模板
                if (StringUtils.isEmpty(userContentRequest.getContent())) {
                    log.warn("用户消息内容为空, sessionId: {}", sessionId);
                    WebSocketUtils.sendMessage(session, "消息内容不能为空");
                    return;
                }
                String content = userContentRequest.getContent();
                String title = content.length() > 20 ? content.substring(0, 20) + "..." : content;

                // 去除换行符等特殊字符
                title = title.replaceAll("[\\r\\n]", " ").trim();

                conversation.setConvTitle(title);
                conversation.setConvDesc("由伴学助手创建的新会话");
                conversation.setConvStatus(1);  // 正常状态
                conversation.setConvType(1);    // 伴学类型
                conversation.setLastChatTime(new Date());
                conversation.setUserId(user.getUser().getUserId());
                conversation.setCreateTime(new Date());
                conversation.setCreateBy(user.getUser().getUserId());
                conversation.setUpdateTime(new Date());
                conversation.setUpdateBy(user.getUser().getUserId());

                KgChatConversationService conversationService = SpringUtils.getBean(KgChatConversationService.class);
                boolean saved = conversationService.save(conversation);
                if (!saved) {
                    log.error("创建新会话失败, userId: {}", user.getUser().getUserId());
                    WebSocketUtils.sendMessage(session, "创建新会话失败");
                    return;
                }

                // 获取自动生成的会话ID
                Long newConvId = conversation.getConvId();
                log.info("新会话创建成功, convId: {}", newConvId);

                // 更新请求中的会话ID
                userContentRequest.setConvId(String.valueOf(newConvId));
            }

            
            AgentServiceIntegration agent = SpringUtils.getBean(AgentServiceIntegration.class);
            ChatMessageRequestV2 request = buildChatMessageRequest(userContentRequest, user.getUser().getUserId());
            WebSocketEventListener eventListener = new WebSocketEventListener(session);
            
            // 异步处理请求
            agent.stream(request, eventListener);
        } catch (Exception e) {
            log.error("处理WebSocket消息时发生错误, sessionId: {}", sessionId, e);
            WebSocketUtils.sendMessage(session, "服务器处理请求失败");
        }
    }

    /**
     * 从WebSocket会话中提取用户信息
     *
     * @param session WebSocket会话
     * @return 登录用户信息，如果未认证则返回null
     */
    private LoginUser extractLoginUser(WebSocketSession session) {
        Principal principal = session.getPrincipal();
        if (principal instanceof UsernamePasswordAuthenticationToken) {
            UsernamePasswordAuthenticationToken authToken = (UsernamePasswordAuthenticationToken) principal;
            return (LoginUser) authToken.getPrincipal();
        }
        return null;
    }

    /**
     * 构建聊天消息请求
     * 
     * @param userContentRequest 用户内容请求
     * @param userId 用户ID
     * @return 聊天消息请求
     */
    private ChatMessageRequestV2 buildChatMessageRequest(ChatAiContentRequest userContentRequest, Long userId) {
        return ChatMessageRequestV2.builder()
                .userId(String.valueOf(userId))
                .convId(userContentRequest.getConvId())
                .content(userContentRequest.getContent())
                .role(ChatMessageRequestV2.Role.USER)
                .messageId(String.valueOf(System.currentTimeMillis()))
                .type(userContentRequest.getType())
                .build();
    }


    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        log.info("[receive binary message] sessionId: {} , message:{}", session.getId(), message.getPayload());
        super.handleBinaryMessage(session, message);
    }

    /**
     * 心跳监测的回复
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        WebSocketUtils.sendPongMessage(session);
    }

    /**
     * 连接出错时
     *
     * @param session
     * @param exception
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("[transport error] sessionId: {} , exception:{}", session.getId(), exception.getMessage());
        if (session.isOpen()) {
            session.close(CloseStatus.SERVER_ERROR);
        }
    }

    /**
     * 连接关闭后
     *
     * @param session
     * @param status
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        WebSocketSessionHolder.removeSession(session.getId());
    }

    /**
     * 指示处理程序是否支持接收部分消息
     *
     * @return 如果支持接收部分消息，则返回true；否则返回false
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
